#coding:utf-8

import pandas as pd
import numpy as np
from transformer_model import train_transformer_model, predict_with_transformer_model

def example_usage():
    """
    Transformer算法使用示例
    """
    print("=== Transformer算法使用示例 ===\n")
    
    # 1. 创建示例数据
    print("1. 创建示例数据...")
    np.random.seed(42)
    n_samples = 500
    
    # 生成特征数据
    df_X = pd.DataFrame({
        'feature1': np.random.randn(n_samples).cumsum(),
        'feature2': np.random.randn(n_samples) * 0.5 + 2,
        'feature3': np.sin(np.arange(n_samples) * 0.05) + np.random.randn(n_samples) * 0.1,
        'feature4': np.cos(np.arange(n_samples) * 0.03) * 2 + np.random.randn(n_samples) * 0.15
    })
    
    # 生成目标数据（带有一些模式）
    df_y = pd.DataFrame({
        'target': (df_X['feature1'] * 0.4 + 
                  df_X['feature2'] * 0.2 + 
                  df_X['feature3'] * 0.3 + 
                  df_X['feature4'] * 0.1 + 
                  np.random.randn(n_samples) * 0.2)
    })
    
    print(f"特征数据形状: {df_X.shape}")
    print(f"目标数据形状: {df_y.shape}")
    print(f"特征列名: {list(df_X.columns)}")
    print(f"目标列名: {list(df_y.columns)}")
    
    # 2. 训练模型
    print("\n2. 训练Transformer模型...")
    model = train_transformer_model(
        df_X=df_X,
        df_y=df_y,
        seq_length=30,      # 序列长度
        epochs=100,         # 训练轮数
        save_path='my_transformer_model.pkl'
    )
    print("模型训练完成并保存为 'my_transformer_model.pkl'")
    
    # 3. 使用模型进行预测
    print("\n3. 使用训练好的模型进行预测...")
    
    # 创建新的测试数据
    test_data = pd.DataFrame({
        'feature1': np.random.randn(50).cumsum() + df_X['feature1'].iloc[-1],
        'feature2': np.random.randn(50) * 0.5 + 2,
        'feature3': np.sin(np.arange(50) * 0.05) + np.random.randn(50) * 0.1,
        'feature4': np.cos(np.arange(50) * 0.03) * 2 + np.random.randn(50) * 0.15
    })
    
    # 进行预测
    predictions = predict_with_transformer_model(
        df_X=test_data,
        model_path='my_transformer_model.pkl'
    )
    
    print(f"输入测试数据行数: {len(test_data)}")
    print(f"预测结果数量: {len(predictions)}")
    print(f"预测结果前10个值: {predictions[:10]}")
    print(f"预测结果统计信息:")
    print(f"  最小值: {min(predictions):.4f}")
    print(f"  最大值: {max(predictions):.4f}")
    print(f"  平均值: {np.mean(predictions):.4f}")
    print(f"  标准差: {np.std(predictions):.4f}")
    
    # 4. 验证预测结果长度一致性
    print(f"\n4. 验证结果:")
    print(f"输入数据行数: {len(test_data)}")
    print(f"预测结果数量: {len(predictions)}")
    print(f"长度一致性: {'✓ 通过' if len(predictions) == len(test_data) else '✗ 失败'}")
    
    print("\n=== 使用示例完成 ===")


def simple_example():
    """
    简化使用示例，展示核心功能
    """
    print("=== 简化使用示例 ===\n")
    
    # 创建简单的数据
    np.random.seed(123)
    df_X = pd.DataFrame({
        'x1': range(100),
        'x2': np.sin(range(100)) * 0.5 + np.random.randn(100) * 0.1
    })
    df_y = pd.DataFrame({
        'y': df_X['x1'] * 0.5 + df_X['x2'] * 2 + np.random.randn(100) * 0.5
    })
    
    print("步骤1: 训练模型...")
    # 训练模型并保存
    model = train_transformer_model(df_X, df_y, seq_length=10, epochs=50, 
                                   save_path='simple_model.pkl')
    
    print("步骤2: 进行预测...")
    # 创建新数据用于预测
    new_data = pd.DataFrame({
        'x1': range(20, 40),
        'x2': np.sin(range(20, 40)) * 0.5 + np.random.randn(20) * 0.1
    })
    
    # 预测
    results = predict_with_transformer_model(new_data, 'simple_model.pkl')
    
    print(f"预测结果: {results}")
    print(f"结果数量: {len(results)} (应该等于输入行数: {len(new_data)})")
    print("=== 简化示例完成 ===\n")


if __name__ == "__main__":
    # 运行完整示例
    example_usage()
    
    # 运行简化示例
    simple_example()